from torchvision import datasets, transforms
from torch.utils.data import DataLoader

def get_data_loader(data_path):
    # Data pipeline
    # 数据预处理
    transform = transforms.Compose(
        [
            transforms.ToTensor(),  # 转换为Tensor并归一化到[0,1]
            transforms.Normalize((0.1307,), (0.3081,)),  # MNIST数据集的全局均值和标准差
        ]
    )

    """ 
        我们在这里指定了数据的路径，避免直接生成在当前路径下。
    """

    # 加载MNIST数据集
    train_dataset = datasets.MNIST(
        data_path, train=True, download=False, transform=transform
    )
    test_dataset = datasets.MNIST(data_path, train=False, transform=transform)

    # 创建数据加载器
    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=1000)
    return train_loader, test_loader
